from genericpath import exists
import os, json, shutil
import os.path as osp

from collections import defaultdict
from flask import Flask
from flask import request

data_root = "/home/xc/work/code/paddle/train_data/det/car/images"
app = Flask(__name__)

# 递归获取文件夹下所有符合条件的文件路径.
def get_all_filepath(root_dir_path, extension_tag=[], recursive = True):
    """
    Args:
        root_dir_path: 文件根目录.
        extension_tag: 文件后缀名组成的列表, Example: ["jpg", "png", "bmp"]
        如果为空 则获取所有
    Returns:
        特定文件夹下所有符合条件的文件路径.
    """
    def check_endswith(new_path, extension_tag):
        if len(extension_tag) == 0:
            return True
        for s in extension_tag:
            if new_path.lower().endswith(s):
                return True
        return False

    paths = []
    for f in os.listdir(root_dir_path):
        new_path = root_dir_path + os.sep + f
        if os.path.isfile(new_path) and check_endswith(new_path,  extension_tag):
            paths.append(new_path)
        elif os.path.isdir(new_path) and recursive:
            temp_list = get_all_filepath(new_path, extension_tag)
            paths.extend(temp_list)
    return paths

# 递归获取文件夹下所有最后一层子文件夹.
def get_last_dir(root_dir_path, recursive_level = -1):
    """
    Args:
        root_dir_path: 文件夹根目录.
    Returns:
        特定文件夹下所有的文件夹路径.
    """
    paths = []
    for f in os.listdir(root_dir_path):
        new_path = root_dir_path + os.sep + f
        if osp.isdir(new_path):
            if recursive_level > 1 or recursive_level < 0:
                ps = get_last_dir(new_path, recursive_level-1)
                paths.extend(ps)
            else:
                paths.append(new_path)
    if len(paths) == 0:
        paths.append(root_dir_path)
    return paths

def get_file_info(data_dir, data_root):
    ret_data = {}
    if not osp.exists(data_dir):
        return ret_data
    
    for f in get_all_filepath(data_dir):
        size = osp.getsize(f)
        ret_data[f.replace(data_root+os.sep, "")] = size
    return ret_data

@app.route('/', methods=['GET', 'POST'])
def home():
    return '<h1>Home</h1>'

@app.route('/get_root', methods=['GET'])
def get_svr_root():
    return json.dumps({"ret":0, "svr_root": data_root})

@app.route('/create_dir', methods=['POST'])
def create_dir():
    data = request.get_json()
    dirs = data["dirs"]
    print("dirs: ", dirs)
    if type(dirs) != list: 
        print("type error...")
        return json.dumps({"ret":100})
    for d in dirs:
        os.makedirs(d, exist_ok=True)
    return json.dumps({"ret":0})

@app.route('/remove_data', methods=['POST'])
def remove_svr_data():
    data = request.get_json()
    print("data: ", data)
    if type(data) != dict: 
        print("type error...", type(data), " != dict")
        return json.dumps({"ret":100})
    remove_data = defaultdict(list)
    for k in data.keys():
        full_path = osp.join(data_root, k)
        remove_data[osp.dirname(full_path)].append(full_path)
    for k,v in remove_data.items():
        total_files = get_all_filepath(k)
        if len(total_files) == len(v):
            print("rmtree: ", k)
            shutil.rmtree(k)
        else:
            print("remove: ", v)
            for vv in v: os.remove(vv)

    return json.dumps({"ret":0})

@app.route('/get_info', methods=['POST'])
def get_all_info():
    query_dir = request.form["query_dir"]
    data_dir = osp.join(data_root, query_dir)
    print("data_dir: ", data_dir)

    ret_data = get_file_info(data_dir, data_root)
    return json.dumps({"ret":0, "data": ret_data})


if __name__ == '__main__':
    app.run("0.0.0.0", 9966)
